"""
五子棋 DQN 系统 -
支持训练模式和人机对弈模式，共用配置系统
"""

import argparse
import os
import sys
from datetime import datetime

from rainforeLearn.gomoku.v2.agents.dqn_agent import GomokuDQNAgent
from rainforeLearn.gomoku.v2.configs.config import GomokuDQNConfig
from rainforeLearn.gomoku.v2.train.self_play_trainer import SelfPlayTrainer
from rainforeLearn.gomoku.v2.gui.gomoku_gui import GomokuGameGUI


class GomokuSystemMain:
    """五子棋 DQN 系统主类"""

    def __init__(self):
        """初始化系统"""
        self.config = None
        self.args = None

    def create_config(self, args):
        """根据命令行参数创建配置"""
        config = GomokuDQNConfig(
            # 棋盘设置
            board_size=args.board_size,

            # 网络配置
            network_config=args.network_config,

            # 训练参数
            replay_buffer_size=args.replay_buffer_size,
            training_stats_batch_freq=args.training_stats_batch_freq,
            learning_rate=args.learning_rate,
            gamma=args.gamma,
            epsilon_start=args.epsilon_start,
            epsilon_end=args.epsilon_end,
            epsilon_decay=args.epsilon_decay,

            # MCTS配置
            use_mcts_sampling=args.use_mcts,
            mcts_update_freq=args.mcts_freq,
            mcts_simulations=args.mcts_simulations,

            # 优先级回放
            use_prioritized_replay=args.use_prioritized_replay,

            # 奖励系统
            enable_win_rate_rewards=args.enable_win_rate_rewards,
            enable_batch_penalty=args.enable_batch_penalty,

            # 文件路径配置
            checkpoint_dir=args.checkpoint_dir,

            # 训练频率配置
            checkpoint_freq=args.checkpoint_freq,
            evaluation_freq=args.evaluation_freq,
            evaluation_games=args.evaluation_games,
            analysis_save_interval=args.analysis_interval,
            log_save_freq=args.log_save_freq,

            # 其他配置
            enable_advanced_analysis=args.enable_analysis,
            monitor_training=args.monitor_training,
        )

        return config

    def train_mode(self, config, args):
        """训练模式"""
        print("🔥 启动五子棋AI训练模式")
        print("=" * 60)

        # 创建智能体
        agent = GomokuDQNAgent(config)

        # 创建训练器
        trainer = SelfPlayTrainer(agent, config)

        # 打印配置信息
        print(f"🎯 训练配置:")
        print(f"  • 棋盘大小: {config.board_size}x{config.board_size}")
        print(f"  • 网络配置: {config.network_config}")
        print(f"  • 训练回合: {args.episodes}")
        print(f"  • 检查点恢复目录: {config.checkpoint_dir}")
        print(f"  • MCTS采样: {'启用' if config.use_mcts_sampling else '禁用'}")
        print(f"  • 优先级回放: {'启用' if config.use_prioritized_replay else '禁用'}")
        print(f"  • 奖励系统: {'启用' if config.enable_win_rate_rewards else '禁用'}")
        print(f"  • 自动保存: {'启用' if config.auto_save else '禁用'}")
        if config.use_expert_data:
            expert_games_count = len(agent.expert_data_loader.expert_games) if agent.expert_data_loader else 0
            print(f"  • 专家棋谱: 启用 ({expert_games_count} 局, 比例: {config.expert_data_ratio:.2f})")
        else:
            print(f"  • 专家棋谱: 禁用")
        if args.checkpoint_filepath:
            print(f"  • 恢复训练: {args.checkpoint_filepath}")

        print("=" * 60)
        # 开始训练
        if args.checkpoint_filepath:
            trainer.train(
                num_episodes=args.episodes,
                checkpoint_filepath=args.checkpoint_filepath
            )
        else:
            trainer.train(num_episodes=args.episodes)

        print("\n🎉 训练完成!")

    def play_mode(self, config, args):
        """游戏模式"""
        print("🎮 启动五子棋人机对弈模式")
        print("=" * 60)

        # 确定模型路径
        model_path = args.model_path
        if model_path is None:
            # 尝试使用默认的最佳模型
            default_model = os.path.join(config.checkpoint_dir, "best_model.pth")
            if os.path.exists(default_model):
                model_path = default_model
                print(f"🤖 使用默认AI模型: {model_path}")
            else:
                print("⚠️ 未指定AI模型，将启动双人对弈模式")
        else:
            if not os.path.exists(model_path):
                print(f"❌ 指定的模型文件不存在: {model_path}")
                print("🎲 将启动双人对弈模式")
                model_path = None
            else:
                print(f"🤖 加载AI模型: {model_path}")

        # 打印游戏配置
        print(f"🎯 游戏配置:")
        print(f"  • 棋盘大小: {config.board_size}x{config.board_size}")
        print(f"  • AI先手: {'是' if args.ai_first else '否'}")
        print(f"  • AI思考时间: {args.ai_thinking_time}秒")
        print(f"  • 显示AI思考: {'是' if args.show_ai_thinking else '否'}")
        print("=" * 60)

        try:
            # 创建游戏
            game = GomokuGameGUI(
                config=config,
                model_path=model_path,
                ai_plays_first=args.ai_first
            )

            # 设置AI参数
            if hasattr(game, 'ai_thinking_time'):
                game.ai_thinking_time = args.ai_thinking_time
            if hasattr(game, 'show_ai_thinking'):
                game.show_ai_thinking = args.show_ai_thinking

            if game.gui_enabled:
                print("🎮 游戏控制:")
                print("  • 鼠标点击: 落子")
                print("  • R键: 重新开始")
                print("  • S键: 保存AI模型 (如果启用AI)")
                print("  • ESC键: 退出游戏")
                print("=" * 60)

                # 开始游戏
                game.play()
            else:
                print("❌ GUI无法初始化，请检查Pygame安装")

        except KeyboardInterrupt:
            print("\n👋 游戏被用户中断")
        except Exception as e:
            print(f"\n❌ 游戏过程中出现错误: {e}")
            import traceback
            traceback.print_exc()

    def setup_argument_parser(self):
        """设置命令行参数解析器"""
        parser = argparse.ArgumentParser(
            description='五子棋 DQN 系统 - 统一主程序',
            formatter_class=argparse.ArgumentDefaultsHelpFormatter
        )

        # === 基本模式选择 ===
        parser.add_argument('--mode', choices=['train', 'play'], default='play',
                            help='运行模式: train=训练AI, play=人机对弈')

        # === 游戏基础配置 ===
        parser.add_argument('--board_size', type=int, default=15,
                            help='棋盘大小')

        # === 网络和学习配置 ===
        parser.add_argument('--network_config', choices=['small', 'medium', 'large'],
                            default='small', help='网络规模配置')
        parser.add_argument('--replay_buffer_size', type=int, default=50000,
                            help='经验回放缓冲区大小')
        parser.add_argument('--batch_size', type=int, default=64,
                            help='批处理大小')
        parser.add_argument('--learning_rate', type=float, default=0.001,
                            help='学习率')
        parser.add_argument('--gamma', type=float, default=0.99,
                            help='折扣因子')
        parser.add_argument('--epsilon_start', type=float, default=1.0,
                            help='探索率初始值')
        parser.add_argument('--epsilon_end', type=float, default=0.01,
                            help='探索率最终值')
        parser.add_argument('--epsilon_decay', type=int, default=10000,
                            help='探索率衰减步数')

        # === MCTS配置 ===
        parser.add_argument('--use_mcts', action='store_true', default=True,
                            help='启用MCTS采样')
        parser.add_argument('--mcts_freq', type=int, default=10,
                            help='MCTS更新频率')
        parser.add_argument('--mcts_simulations', type=int, default=100,
                            help='MCTS模拟次数')

        # === 高级训练功能 ===
        parser.add_argument('--use_prioritized_replay', action='store_true', default=True,
                            help='启用优先级经验回放')
        parser.add_argument('--enable_win_rate_rewards', action='store_true', default=True,
                            help='启用胜率奖励系统')
        parser.add_argument('--enable_batch_penalty', action='store_true',default=True,
                            help='启用批次惩罚系统')

        # === 文件和目录配置 ===
        parser.add_argument('--checkpoint_dir', type=str, default='checkpoints',
                            help='检查点文件目录')
        parser.add_argument('--checkpoint_filepath', type=str,
                            help='检查点文名')

        # === 训练模式专用参数 ===
        train_group = parser.add_argument_group('训练模式参数')
        train_group.add_argument('--episodes', type=int, default=1000,
                                 help='训练回合数')

        train_group.add_argument('--checkpoint_freq', type=int, default=100,
                                 help='检查点保存频率')
        train_group.add_argument('--evaluation_freq', type=int, default=50,
                                 help='评估频率')
        train_group.add_argument('--evaluation_games', type=int, default=50,
                                 help='评估游戏数量')
        train_group.add_argument('--training_stats_batch_freq', type=int, default=10,
                                 help='统计打印频率')
        train_group.add_argument('--analysis_interval', type=int, default=100,
                                 help='分析图表保存间隔')
        train_group.add_argument('--log_save_freq', type=int, default=50,
                                 help='日志保存频率')
        train_group.add_argument('--enable_analysis', action='store_true', default=True,
                                 help='启用高级分析')
        train_group.add_argument('--monitor_training', action='store_true', default=True,
                                 help='启用训练监控')

        # === 游戏模式专用参数 ===
        game_group = parser.add_argument_group('游戏模式参数')
        game_group.add_argument('--model_path', type=str, default=None,
                                help='AI模型文件路径')
        game_group.add_argument('--ai_first', action='store_true', default=True,
                                help='AI先手(黑棋)')
        game_group.add_argument('--ai_thinking_time', type=float, default=0.5,
                                help='AI思考时间(秒)')
        game_group.add_argument('--show_ai_thinking', action='store_true', default=True,
                                help='显示AI思考状态')

        return parser

    def run(self, args=None):
        """主运行方法"""
        # 设置参数解析器
        parser = self.setup_argument_parser()

        # 解析参数
        if args is None:
            self.args = parser.parse_args()
        else:
            self.args = parser.parse_args(args)

        # 创建配置
        self.config = self.create_config(self.args)

        # 确保检查点目录存在
        os.makedirs(self.config.checkpoint_dir, exist_ok=True)

        # 打印启动信息
        print("🚀 五子棋 DQN 系统")
        print(f"⏰ 启动时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        print(f"🎯 运行模式: {self.args.mode.upper()}")
        print()

        # 根据模式启动相应功能
        if self.args.mode == 'train':
            self.train_mode(self.config, self.args)
        elif self.args.mode == 'play':
            self.play_mode(self.config, self.args)

        print("\n🏁 程序结束")


def main():
    """主函数 - 保持向后兼容"""
    system = GomokuSystemMain()
    system.run()


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n👋 程序被用户中断")
        sys.exit(0)
    except Exception as e:
        print(f"\n❌ 程序运行错误: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)